﻿import threading
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from enum import Enum
import numpy as np
from numpy import ndarray, uint8, uint16, uint32, uint64, float32, float64


@dataclass
class GlobalMemory:
    def __init__(self, max_size: int = 1024**3):
        self.__max: int = max_size
        self.__mem: Dict[int, List[threading.Lock, ndarray[uint8]]] = {}

    def read(self, addr: int, size: str) -> int | float:
        if addr+size > self.__max:
            raise Exception(
                f"Memory access out of bounds: {addr}+{size}>{self.__max}")
        pfn, offset = divmod(addr, 4096)
        if pfn not in self.__mem:
            raise Exception(f"Memory access to unmapped page: {pfn}")
        lock, data = self.__mem[pfn]
        with lock:
            match size:
                case 'byte':
                    return data[offset]
                case 'word':
                    return data[offset:offset+2].view(uint16)[0]
                case 'dword':
                    return data[offset:offset+4].view(uint32)[0]
                case 'qword':
                    return data[offset:offset+8].view(uint64)[0]
                case 'float':
                    return data[offset:offset+4].view(float32)[0]
                case 'double':
                    return data[offset:offset+8].view(float64)[0]

    def write(self, addr: int, size: str, data: int | float):
        if addr+size > self.__max:
            raise Exception(
                f"Memory access out of bounds: {addr}+{size}>{self.__max}")
        pfn, offset = divmod(addr, 4096)
        if pfn not in self.__mem:
            self.__mem[pfn] = (threading.Lock(), bytearray(4096))
        lock = self.__mem[pfn][0]
        with lock:
            # 直接更新字典中的bytes对象
            match size:
                case 'byte':
                    self.__mem[pfn][1][offset].view(uint8)[0] = data
                case 'word':
                    self.__mem[pfn][1][offset:offset+2].view(uint16)[0] = data
                case 'dword':
                    self.__mem[pfn][1][offset:offset+4].view(uint32)[0] = data
                case 'qword':
                    self.__mem[pfn][1][offset:offset+8].view(uint64)[0] = data
                case 'float':
                    self.__mem[pfn][1][offset:offset+4].view(float32)[0] = data
                case 'double':
                    self.__mem[pfn][1][offset:offset+8].view(float64)[0] = data

    def get(self, addr: int) -> ndarray[uint8]:
        pfn, offset = divmod(addr, 4096)
        if pfn not in self.__mem:
            raise Exception(f"Memory access to unmapped page: {pfn}")
        lock, data = self.__mem[pfn]
        with lock:
            return data[offset:offset+64].view(uint8)

    def put(self, addr: int, data: ndarray[uint8]) -> None:
        pfn, offset = divmod(addr, 4096)
        if pfn not in self.__mem:
            self.__mem[pfn] = (threading.Lock(), data)
        lock = self.__mem[pfn][0]
        with lock:
            self.__mem[pfn][1][offset:offset+64].view(uint8)[:] = data

    def __setitem__(self, addr: int, data: ndarray[uint8]) -> None:
        pfn, offset = divmod(addr, 4096)
        if pfn not in self.__mem:
            self.__mem[pfn] = (
                threading.Lock(), np.zeros(shape=4096, dtype=uint8))
        lock = self.__mem[pfn][0]
        if (offset+len(data)) >= 4096:
            raise Exception(
                f"Memory access across pages at {addr},length:{len(data)}")
        with lock:
            self.__mem[pfn][1][offset:offset+len(data)] = data

    def __getitem__(self, addr: int) -> ndarray[uint8]:
        pfn, offset = divmod(addr, 4096)
        if pfn not in self.__mem:
            raise Exception(f"Memory access to unmapped page: {pfn}")
        lock, data = self.__mem[pfn]
        with lock:
            return data
